//
// Created by lidongwei on 2023/6/21.
//

#include "infer.h"
#include <fstream>

//bool __check_cuda_runtime(cudaError_t code, const char *op, const char *file, int line) {
//    if (code != cudaSuccess) {
//        const char *err_name = cudaGetErrorName(code);
//        const char *err_message = cudaGetErrorString(code);
//        printf("runtime error %s:%d  %s failed. \n  code = %s, message = %s\n", file, line, op, err_name, err_message);
//        return false;
//    }
//    return true;
//}

/////////////////////////////////打印推理日志的函数/////////////////////
inline const char *severity_string(nvinfer1::ILogger::Severity t) {
    switch (t) {
        case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
            return "internal_error";
        case nvinfer1::ILogger::Severity::kERROR:
            return "error";
        case nvinfer1::ILogger::Severity::kWARNING:
            return "warning";
        case nvinfer1::ILogger::Severity::kINFO:
            return "info";
        case nvinfer1::ILogger::Severity::kVERBOSE:
            return "verbose";
        default:
            return "unknow";
    }
}

class TRTLogger : public nvinfer1::ILogger {
public:
    virtual void log(Severity severity, nvinfer1::AsciiChar const *msg) noexcept override {
        if (severity <= Severity::kINFO) {
            // 打印带颜色的字符，格式如下：
            // printf("\033[47;33m打印的文本\033[0m");
            // 其中 \033[ 是起始标记
            //      47    是背景颜色
            //      ;     分隔符
            //      33    文字颜色
            //      m     开始标记结束
            //      \033[0m 是终止标记
            // 其中背景颜色或者文字颜色可不写
            // 部分颜色代码 https://blog.csdn.net/ericbar/article/details/79652086
            if (severity == Severity::kWARNING) {
                printf("\033[33m%s: %s\033[0m\n", severity_string(severity), msg);
            } else if (severity <= Severity::kERROR) {
                printf("\033[31m%s: %s\033[0m\n", severity_string(severity), msg);
            } else {
                printf("%s: %s\n", severity_string(severity), msg);
            }
        }
    }
} logger;

////////////////////////////////推理模块//////////////////////////////
std::vector<unsigned char> trt_infer::Infer::load_file(const std::string &file) {//定义读取engine文件的函数
    std::ifstream in(file, std::ios::in | std::ios::binary);
    if (!in.is_open())
        return {};

    in.seekg(0, std::ios::end);
    size_t length = in.tellg();

    std::vector<uint8_t> data;
    if (length > 0) {
        in.seekg(0, std::ios::beg);
        data.resize(length);

        in.read((char *) &data[0], length);
    }
    in.close();
    return data;
}

//std::vector<float> trt_infer::Infer::infer(std::vector<unsigned char> engine_model, std::vector<float> input_data,int batch,int channel,int input_width,int input_height,int output_size) {
//    TRTLogger logger;//日志
//    nvinfer1::IRuntime *runtime = nvinfer1::createInferRuntime(logger);//导入日志配置
//    nvinfer1::ICudaEngine *engine = runtime->deserializeCudaEngine(engine_model.data(), engine_model.size());//模型反序列化
//    if (engine == nullptr) {
//        printf("Deserialize cuda engine failed.\n");//反序列化cuda引擎失败
//        runtime->destroy();
//        std::vector<float> re;
//        return re;
//    }
//
//    nvinfer1::IExecutionContext *executionContext = engine->createExecutionContext();//创建上下文
//    cudaStream_t stream = nullptr;
//    cudaStreamCreate(&stream);
//
//    float *input_data_device= nullptr;
//    std::vector<float> output_data_host(output_size);
//    float *output_data_device= nullptr;
//    checkRuntime(cudaMalloc(&input_data_device,sizeof(float)*input_data.size()));
//    checkRuntime(cudaMalloc(&output_data_device,sizeof(float)*output_size));
//
//    auto start=clock();
//    auto input_dims = executionContext->getBindingDimensions(0);
//    input_dims.d[0] = batch;
//    checkRuntime(cudaMemcpy(input_data_device,input_data.data(),sizeof(float)*input_data.size(),cudaMemcpyHostToDevice));
//
//    //设置推理式，使用的数据输入大小
//    executionContext->setBindingDimensions(0,input_dims);
//    float *bindings[]={input_data_device,output_data_device};
//    for(int i =1;i<5;i++){
//        bool success=executionContext->enqueueV2((void**)bindings,stream, nullptr);//推理
//        checkRuntime(cudaStreamSynchronize(stream));
//    }
//    checkRuntime(cudaMemcpy(output_data_host.data(),output_data_device,sizeof(float)*output_size,cudaMemcpyDeviceToHost));
////    checkRuntime(cudaStreamSynchronize(stream));
//    auto end=clock();
//
//    double diff_time=(double)(end-start)/CLOCKS_PER_SEC*1000;
//    printf("推理成功：输出长度：%d x %d ---- 耗时：%.1f ms\n",batch,(output_data_host.size()/batch),diff_time);
//
//    checkRuntime(cudaFree(output_data_device));
//    checkRuntime(cudaFree(input_data_device));
//
//    return output_data_host;
//}



nvinfer1::IExecutionContext * trt_infer::Infer::create_engine(std::vector<unsigned char> engine_model,int batch) {
    TRTLogger logger;//日志
    nvinfer1::IRuntime *runtime = nvinfer1::createInferRuntime(logger);//导入日志配置
    nvinfer1::ICudaEngine *engine = runtime->deserializeCudaEngine(engine_model.data(), engine_model.size());//模型反序列化
    if (engine == nullptr) {
        printf("Deserialize cuda engine failed.\n");//反序列化cuda引擎失败
        runtime->destroy();
        nvinfer1::IExecutionContext *executionContext= nullptr;
        return executionContext;
    }

    nvinfer1::IExecutionContext *executionContext = engine->createExecutionContext();//创建上下文

    auto input_dims = executionContext->getBindingDimensions(0);
    input_dims.d[0] = batch;

    //设置推理式，使用的数据输入大小
    executionContext->setBindingDimensions(0,input_dims);
    return executionContext;

}


std::vector<float> trt_infer::Infer::infer(nvinfer1::IExecutionContext * executionContext, std::vector<float> input_data,int batch,int channel,int input_width,int input_height,int output_size) {

    auto start=clock();
    cudaStream_t stream = nullptr;
    cudaStreamCreate(&stream);

    float *input_data_device= nullptr;
    std::vector<float> output_data_host(output_size);
    float *output_data_device= nullptr;
    checkRuntime(cudaMalloc(&input_data_device,sizeof(float)*input_data.size()));
    checkRuntime(cudaMalloc(&output_data_device,sizeof(float)*output_size));

    checkRuntime(cudaMemcpy(input_data_device,input_data.data(),sizeof(float)*input_data.size(),cudaMemcpyHostToDevice));

    //推理数据设置
    float *bindings[]={input_data_device,output_data_device};
    bool success=executionContext->enqueueV2((void**)bindings,stream, nullptr);//推理
    checkRuntime(cudaStreamSynchronize(stream));

    checkRuntime(cudaMemcpy(output_data_host.data(),output_data_device,sizeof(float)*output_size,cudaMemcpyDeviceToHost));
    auto end=clock();

    double diff_time=(double)(end-start)/CLOCKS_PER_SEC*1000;
    printf("推理成功：输出长度：%d x %d ---- 耗时：%.1f ms\n",batch,(output_data_host.size()/batch),diff_time);

    checkRuntime(cudaFree(output_data_device));
    checkRuntime(cudaFree(input_data_device));

    return output_data_host;
}